-
Notifications
You must be signed in to change notification settings - Fork 3.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Unify name mangling in TVM #12066
Unify name mangling in TVM #12066
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Given that we can recreate GlobalVarSupply from the current IRModule, and we don't need global var renaming for most of the passes(that only do local edits of functions).
It might make sense to create helper method to generate GlobalVarSupply rather than maintaining as a member of IRModule when needed, to keep the main IRModule data structure simple.
Doing so would also simplify the serialization part of IRModule, right now the unordered map data structure is not serialized, as a result save_json/load_json of IRModule will no longer work if we make it as a member.
include/tvm/ir/module.h
Outdated
@@ -64,6 +65,8 @@ class IRModuleNode : public Object { | |||
/* \brief Additional attributes storing meta-data about the module. */ | |||
DictAttrs attrs; | |||
|
|||
GlobalVarSupply global_var_supply; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Given that we can recreate GlobalVarSupply from the current IRModule, and we don't need global var renaming for most of the passes(that only do local edits of functions).
It might make sense to create helper method to generate GlobalVarSupply rather than maintaining as a member of IRModule when needed, to keep the main IRModule data structure simple.
Doing so would also simplify the serialization part of IRModule, right now the unordered map data structure is not serialized, as a result save_json/load_json of IRModule will no longer work if we make it as a member.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Given that we can recreate GlobalVarSupply from the current IRModule
I'm not sure this is so trivial. A couple of cases to think through:
- right now we still lower per-backend separately, so the NameSupply needs to track names across a collection of IRModules in the compiler
- Suppose we needed to reserve an extern name not mentioned in the IR anywhere. We would surely need to serialize something to track that.
The first case here seems to be the one most at tension with the idea of making NameSupply idempotent. I agree that NameSupply isn't necessarily attached to any given IRModule now, so having a way to access it should mainly be a convenience method and not participate in IRModule serialization at least right now. I'm not so sure it means that it's valid to serialize an IRModule and destroy a TVM instance, then load it back in another TVM instance and expect the NameSupply to be the same.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I agree @areusch , both of your pts seems to point that we should bring NameSupply outside IRModule, and have additional functionalities to re-create from multple (collection) of IRModules
Alternatively, we can make name supply an Even in this alternative case, we need to work carefully with cases where we copy on write the IRModule. In nature NameSupply is a mutable object while IRModule is moving towards more immutable style. e.g. if we copy an IRModule and mutate it once, we need to understand the API impact on the original IRModule. Allowing recreation of NameSupply from fresh might be the best way to avoid such problems. So possible options:
Example to demonstrate the branching case def transform_example(moda):
modb = tvm.transform.SomePass1()(moda)
modc = tvm.transform.SomePass2()(moda) In the above example, we can either choose to persist NameSupply during transformation, as a result modb, modc share the same NameSupply, causing bugs since they really are two different pathways. Or we duplicate/recreate the NameSupply in each pass(since we do not know whether or not the NameSupply will be used in another mod), which leads to overhead. This is of course follows the design rationale of keeping things immutable, in this world, likely likely F0 and F1 will have similar effects and recreation when needed(in few times that it is necessary) may not be a bad idea. |
a2fc6db
to
fb27659
Compare
Thank you @tqchen for the review. I also prefer F0, where we always create for immediate use a fresh One possible issue I can see with both F0 and F1 is that usually the module name is used as a prefix to the GlobalVars that belong to the same module. Currently, the module name is not a member of |
fb27659
to
a7a9278
Compare
Thanks @gigiblender Adding module name to IRModule as an optional attr sounds great. BTW, we still need to distinguish the case where the user expect "the name". This is different from cases when a pass wants to allocate a new name that will be used later, but there is no expectation of external reference(as a result the name can be flexible). This happens in cases where user expects an external function name that they can call during runtime. In which case NameSupply can not do anything but to check duplication. |
4ccc004
to
96a7135
Compare
84ae2ef
to
428a598
Compare
c3182af
to
6962c2f
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
thanks @gigiblender, left a round of comments
cc @kparzysz-quic @Mousius @manupa-arm @junrushao1994 @Hzfengsy @tqchen @comaniac please review |
e1177da
to
30c27af
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I reviewed half of it so far.
Make sure that your algorithm can handle this (pseudo code):
get_new_name("name2"); // creates "name2"
get_new_name("name"); // creates "name"
get_new_name("name"); // creates "name1"
get_new_name("name"); // does not create "name2" again!
src/ir/name_supply.cc
Outdated
for (size_t i = 0; i < prefix.size(); ++i) { | ||
if (prefix[i] == '.') prefix[i] = '_'; | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why are you doing this replacement?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think the dot operator is used to select a member of a class/struct and therefore should be mangled in all names.
f4ee874
to
25ca250
Compare
@gigiblender is this one ready for another look? |
Yes. All comments addressed so far. |
25ca250
to
5df20bb
Compare
* [TVM] Update Submodule * [Compatible] Fix apache/tvm#12066 * [Compatible] Fix apache/tvm#12382 Co-authored-by: SubmoduleUpdaterBot <submodule-updater-bot@users.noreply.github.com> Co-authored-by: Cody Yu <comaniac0422@gmail.com>
* Add NameSupply and GlobalVarSupply * Build GlobalVarSupply from IRModules instead of having it attached to an IRModule. * Pass GlobalVarSupply when lowering shape funcs * Partially replace instantiations of GlobalVar with GlobalVarSupply * Construct GlobalVarSupply from IRModule * Add tests for supply * Add documentation for NameSupply and GlobalVarSupply Co-authored-by: Florin-Gabriel Blanaru <fgb@system76-pc.localdomain>
* Add NameSupply and GlobalVarSupply * Build GlobalVarSupply from IRModules instead of having it attached to an IRModule. * Pass GlobalVarSupply when lowering shape funcs * Partially replace instantiations of GlobalVar with GlobalVarSupply * Construct GlobalVarSupply from IRModule * Add tests for supply * Add documentation for NameSupply and GlobalVarSupply Co-authored-by: Florin-Gabriel Blanaru <fgb@system76-pc.localdomain>
This PR is an implementation for this RFC discussion.
Tracking issue: #12181
Forum discussion: https://discuss.tvm.apache.org/t/pre-rfc-name-mangling-in-irmodules/12944/7
It implements a
NameSupply
andGlobalVarSupply
to be used in TVM for obtaining unique names / global vars. It also replaces the previous usages of aname_map
and theGetUniqueName
method with aNameSupply
, orGlobalVarSupply
depending on the case.